import os
import lightning as L
import pandas as pd
import torch
from tlhengine.utils import random_split, get_pickle
from torch.utils.data.dataloader import DataLoader
from datasets import CustomDataset, TestDataset, VideoDataset
from engine import Model
from torchvision import transforms
from compare import compare
from argparse import ArgumentParser

def predict_video_loader(trainer, model, loader):
    model.eval()
    try:
        for videos in loader:
            # print(videos.shape)
            # print(predictions)
            # predictions = model(videos)
            videos = videos.squeeze()
            predictions = model(videos)
            predictions = predictions > 0.5
            predictions = torch.all(predictions, dim=1)
            print(predictions)
    except:
        pass

parser = ArgumentParser()
parser.add_argument('--weights',
                    default='/root/code/tianchi/log/noname/version_38/checkpoints/best_0.7759259343147278.ckpt')
args = parser.parse_args()
test_preprocessing = transforms.Compose([
            transforms.Resize((224, 224)),
            # transforms.RandomHorizontalFlip(),
            # transforms.RandomVerticalFlip(),
            transforms.ToTensor(),
            transforms.Normalize( mean = [0.485, 0.456, 0.406] , std = [0.229, 0.224, 0.225])
        ])
video_test_preprocessing = transforms.Compose([
            transforms.ToPILImage(),
            transforms.Resize((224, 224)),
            # transforms.RandomHorizontalFlip(),
            # transforms.RandomVerticalFlip(),
            transforms.ToTensor(),
            transforms.Normalize( mean = [0.485, 0.456, 0.406] , std = [0.229, 0.224, 0.225])
        ])


test_set = TestDataset("data/test/picture", test_preprocessing)
# test_set = TestDataset("data/test/picture", test_preprocessing)

test_video_set = VideoDataset('/root/code/tianchi/data/test/video', transform=video_test_preprocessing)
test_video_loader = DataLoader(
    test_video_set,
    batch_size= 1 ,
    num_workers=4 
)

test_loader = DataLoader(
    test_set,
    batch_size=32,
    num_workers=4
)
model = Model.load_from_checkpoint(
    args.weights
)

trainer = L.Trainer(max_epochs=100, logger=False)
predictions = trainer.predict(model, dataloaders=test_loader)
# predict_video_loader(trainer, model, test_video_loader)
# model.predci(test_loader)
print("done")
predictions_cat = torch.cat(
    predictions,
)
def convert_to_bit(preds):
    bit_values = []
    for pred in preds:
        # this_pred = pred > 0.5
        this_pred = pred > 0
        
        bit_value = 0
        for i, v in enumerate(this_pred):
            bit_value += v * 2**i
        bit_values.append(bit_value.item())
    return bit_values
bit_values = convert_to_bit(predictions_cat)

filenames = []
for path in test_set.img_paths:
    print(os.path.basename(path))
    filenames.append(os.path.basename(path))
    
df = pd.DataFrame(dict(filename=filenames, result=bit_values))
to_filename= 'test.csv'
df.to_csv(to_filename, index=False)

compare(to_filename)
